import socket
import threading
from select import select


def start_udp_server(address, port):
    udp_server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    udp_server.bind((address, port))
    # read_list = [udp_server.fileno()]
    # write_list = [sys.stdin.fileno()]
    clients = []
    while True:
        data, addr = udp_server.recvfrom(65535)
        if addr not in clients:
            clients.append(addr)

        for c in clients:
            udp_server.sendto(data, c)
        # process data
    udp_server.close()


send_queue = {

}

<<<<<<< HEAD

def start_tcp_server_with_poll(address, port):
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((address, port))
    server.listen(10)
    clients = {}
    # 创建POLLING对象
    p = poll()
    p.register(server.fileno(), POLLIN)
    while True:
        ready_fds = p.poll()
        # ev == POLLIN | POLLOUT
        for fd, ev in ready_fds:
            if ev & POLLIN:
                if fd == server.fileno():
                    client, _ = server.accept()
                    clients[client.fileno()] = client
                    # TODO: add code to generate a data queue
                    for c in clients.keys():
                        p.register(c, POLLOUT)

                    p.register(client.fileno(), POLLIN)
                else:
                    client = clients[fd]
                    try:
                        data = client.recv(65535)
                        if not data:
                            raise ConnectionError
                    except ConnectionError:
                        client.close()
                        p.unregister(fd, POLLIN)
            if ev & POLLOUT:
                pass
=======
sessions = {}

# def start_tcp_server_with_poll(address, port):
#     server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
#     server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
#     server.bind((address, port))
#     server.listen(10)
#     clients = {}
#     # 创建POLLING对象
#     p = poll()
#     p.register(server.fileno(), POLLIN)
#     while True:
#         ready_fds = p.poll()
#         for fd, ev in ready_fds:
#             if ev & POLLIN:
#                 if fd == server.fileno():
#                     client, _ = server.accept()
#                     clients[client.fileno()] = client
#                     # TODO: add code to generate a data queue
#                     # for c in clients.keys():
#                     #     p.register(c, POLLOUT)
#                     p.register(client.fileno(), POLLIN)
#                 else:
#                     client = clients[fd]
#                     try:
#                         data = client.recv(65535)
#                         if not data:
#                             raise ConnectionError
#                     except ConnectionError:
#                         client.close()
#                         p.unregister(fd, POLLIN)
#             if ev & POLLOUT:
#                 pass

write_list = []


def process_data(client, data):
    write_list.append(client)
    if client not in send_queue:
        send_queue[client] = []
    send_queue[client].append(data)
    # 唤醒select线程
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.sendto(b'0', ("127.0.0.1", 8811))
    udp_sock.close()
>>>>>>> bc2540887941e532399a8bc07191889337a4419e


def start_tcp_server_with_select(address, port):
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((address, port))
    server.listen(10)

    # 该UDP套接口用于唤醒等待的select事件
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.bind(("127.0.0.1", 8811))

    read_list = [server, udp_sock]

    while True:
        read_fds, write_fds, _ = select(read_list, write_list, [])
        for fd in read_fds:
            assert isinstance(fd, socket.socket)  
            if fd == server:
                client, addr = fd.accept()
                read_list.append(client)
                print("{} is online".format(addr))
            elif fd == udp_sock:
                _, _ = udp_sock.recvfrom(1)
                continue
            else:
                try:
                    data = fd.recv(65535)
                    if not data:
                        raise ConnectionError
                    try:
                        print('{} sent a message:{}'.format(fd.getpeername(), data.decode("utf-8")))
                    except:
                        pass
                    t = threading.Thread(target=process_data,
                                         args=(fd, data))
                    t.daemon = True
                    t.start()
                except ConnectionError:
                    read_list.remove(fd)
                    fd.close()

        for fd in write_fds:
            if fd not in send_queue:
                write_list.remove(fd)
                continue
            data_queue = send_queue[fd]
            data = data_queue.pop(0)
            fd.send(data)
            if not data_queue:
                del send_queue[fd]
                write_list.remove(fd)


import argparse
import sys

print(sys.argv)

ap = argparse.ArgumentParser()
ap.add_argument("-t", type=str, choices=("tcp", "udp"), help="to use tcp / udp", default="tcp")
ap.add_argument("-H", type=str, required=True, help="the ip address to bind")
ap.add_argument("-p", type=int, required=True, help="the port to bind")

arg = ap.parse_args(sys.argv[1:])
if arg.t == "tcp":
    start_tcp_server_with_select(arg.H, arg.p)
else:
    start_udp_server(arg.H, arg.p)
